from keras.datasets import cifar10
from keras.preprocessing.image import ImageDataGenerator
from keras.models import Sequential
from keras.layers.core import Dense, Dropout, Activation, Flatten
from keras.layers.convolutional import Convolution2D, MaxPooling2D
from keras.optimizers import SGD,adam
from keras.utils import np_utils
from Data_utils.Load import load_train17,load_val17,load_train32,load_val32
import numpy as np
from Data_utils.Crrupt import darken,add_gaussian_noise
from model import CAE_model
import os
from skimage import io,util,exposure
import matplotlib.pyplot as plt
import sys
import keras.backend as K
from Data_utils.Metrics import PSNR
import json
import numpy
"""
def PSNR_loss(I,K):
    batch_size = Kb.eval(Kb.variable(I)).shape[0]
    I = Kb.reshape(I,(int(batch_size),17*17))
    K = Kb.reshape(K,(int(batch_size),17*17))
    MSE = Kb.sum(Kb.square(I-K),axis=1)
    PSNR_ = 20. * Kb.log10(255./Kb.sqrt(MSE))
    PSNR_ = Kb.mean(PSNR_)
    return (1/2.)**PSNR_
"""
def train_model():
    train_imgs, train_names = load_train17(Debug=True)
    val_imgs, val_names = load_val17(Debug=True)
    train_Y = np.asarray(train_imgs)
    train_X = add_gaussian_noise(train_Y)
    val_Y = np.asarray(val_imgs)
    val_X = add_gaussian_noise(val_Y)
    print '{0} train samples'.format(len(train_Y))
    print '{0} val samples'.format(len(val_Y))
    print 'patch size: {0}'.format(train_Y[0].shape[0])

    train_X = np.reshape(train_X,(train_X.shape[0],1,train_X[0].shape[0],train_X[0].shape[1]))
    train_Y = np.reshape(train_Y,(train_Y.shape[0],1,train_Y[0].shape[0],train_Y[0].shape[1]))
    train_X /= 255.
    train_Y /= 255.
    val_X = np.reshape(val_X,(val_X.shape[0],1,val_X[0].shape[0],val_X[0].shape[1]))
    val_Y = np.reshape(val_Y,(val_Y.shape[0],1,val_Y[0].shape[0],val_Y[0].shape[1]))
    val_X /= 255.
    val_Y /= 255.
    model = CAE_model(img_channels=1,img_rows=17,img_cols=17)

    model.compile(optimizer=SGD(lr=0.001,momentum=0.9,nesterov=False),loss='mse')
    hist = model.fit(X=train_X,y=train_Y,nb_epoch=20,validation_data=(val_X,val_Y),batch_size=256,shuffle=True)

    model.save_weights('metadata/darken_weights_mse.h5',overwrite=True)
    training = {
        'loss': hist.history['loss'],
        'val_loss': hist.history['val_loss'],
        'optimizer': model.optimizer.get_config(),
    }
    f = open('metadata/darken_training_mse.json', 'wb')
    meta_json = json.dumps(training, default=lambda o: o.__dict__, indent=4)
    f.write(meta_json)
    f.close()

def inference_darken():
    model = CAE_model(img_channels=1,img_rows=512,img_cols=512)
    model.load_weights('metadata/darken_weights_mse.h5')
    nb_img = 5
    img_folder = '/media/dell/cb552bf1-c649-4cca-8aca-3c24afca817b/dell/wxm/Data/decsai/GM512/inf'
    for img in os.listdir(img_folder):
        if nb_img <= 0:
            break
        nb_img -= 1
        img_path = img_folder + '/' + img
        img_arr = io.imread(img_path)
        img_arr_darked = exposure.adjust_gamma(image=img_arr,gamma=np.random.uniform(2.,5.))
        img_arr_darked = np.reshape(img_arr_darked,(1,1,img_arr_darked.shape[0],img_arr_darked.shape[1]))
        print img_arr_darked.shape
        pred = K.eval(model(K.variable(img_arr_darked/255.)))
        pred = np.asarray(pred)
        print pred.shape
        img_recostructed = pred * 255.

        img_arr_darked = np.reshape(img_arr_darked,(img_arr.shape[0],img_arr.shape[1]))
        img_recostructed = np.reshape(img_recostructed,(img_arr.shape[0],img_arr.shape[1]))
        #img_recostructed = np.clip(img_recostructed,0,255)
        # show images
        plt.subplot(1,3,1)
        io.imshow(img_arr)
        plt.subplot(1,3,2)
        io.imshow(img_arr_darked)
        plt.subplot(1,3,3)
        io.imshow(img_recostructed)
        io.show()
        print 'PSNR: {0}'.format(PSNR(I=img_arr,K=img_recostructed))

def inference_noise():
    model = CAE_model(img_channels=1,img_rows=512,img_cols=512)
    model.load_weights('metadata/noise_weights_mse.h5')
    nb_img = 5
    img_folder = '/media/dell/cb552bf1-c649-4cca-8aca-3c24afca817b/dell/wxm/Data/decsai/GM512/inf'
    for img in os.listdir(img_folder):
        if nb_img <= 0:
            break
        nb_img -= 1
        img_path = img_folder + '/' + img
        img_arr = io.imread(img_path)
        img_arr_noised = util.random_noise(image=img_arr,var=((25/255.)**2)*np.random.uniform(0,1))
        img_arr_noised = np.reshape(img_arr_noised,(1,1,img_arr_noised.shape[0],img_arr_noised.shape[1]))
        print img_arr_noised.shape
        pred = K.eval(model(K.variable(img_arr_noised/255.)))
        pred = np.asarray(pred)
        print pred.shape
        img_recostructed = pred * 255.

        img_arr_noised = np.reshape(img_arr_noised,(img_arr.shape[0],img_arr.shape[1]))
        img_recostructed = np.reshape(img_recostructed,(img_arr.shape[0],img_arr.shape[1]))
        #img_recostructed = np.clip(img_recostructed,0,255)
        # show images
        plt.subplot(1,3,1)
        io.imshow(img_arr)
        plt.subplot(1,3,2)
        io.imshow(img_arr_noised)
        plt.subplot(1,3,3)
        io.imshow(img_recostructed)
        io.show()
        print 'PSNR: {0}'.format(PSNR(I=img_arr,K=img_recostructed))

if __name__ == '__main__':
    inference_darken()